import os
import cv2
import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.callbacks import EarlyStopping
from sklearn.model_selection import train_test_split


def load_dataset(data_dir):
    gestures = ['paper', 'rock', 'scissors']
    X = []
    y = []

    for gesture_idx, gesture in enumerate(gestures):
        gesture_dir = os.path.join(data_dir, gesture)
        for img_name in os.listdir(gesture_dir):
            img_path = os.path.join(gesture_dir, img_name)
            img = cv2.imread(img_path)

            if img is None:
                continue

            img = cv2.resize(img, (64, 64))  # 保持64x64大小
            img = img.astype('float32') / 255.0
            X.append(img)
            y.append(gesture_idx)

    return np.array(X), np.array(y)


def train_model(X, y):
    # 使用原有模型结构，只添加Dropout层
    model = Sequential([
        Conv2D(32, (3, 3), activation='relu', input_shape=(64, 64, 3)),
        MaxPooling2D(2, 2),

        Conv2D(64, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),

        Conv2D(128, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),

        Flatten(),
        Dense(256, activation='relu'),
        Dropout(0.5),  # 添加Dropout防止过拟合
        Dense(3, activation='softmax')
    ])

    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    # 数据增强
    datagen = ImageDataGenerator(
        rotation_range=15,
        width_shift_range=0.1,
        height_shift_range=0.1,
        horizontal_flip=True
    )

    # 划分训练集和验证集
    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

    # 使用早停
    early_stop = EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True)

    # 训练模型
    model.fit(
        datagen.flow(X_train, y_train, batch_size=32),
        steps_per_epoch=len(X_train) // 32,
        epochs=30,
        validation_data=(X_val, y_val),
        callbacks=[early_stop]
    )

    return model


if __name__ == "__main__":
    # 确保模型目录存在
    if not os.path.exists("models"):
        os.makedirs("models")

    # 加载数据集
    X, y = load_dataset("datasets")

    # 训练模型
    model = train_model(X, y)

    # 保存模型 - 使用原有文件名
    model.save("models/gesture_model.h5")
    print("模型训练完成并已保存")